import grpc
import llm_pb2
import llm_pb2_grpc
import argparse

def run(question):
    with grpc.insecure_channel('localhost:50051') as channel:
        stub = llm_pb2_grpc.ChatServiceStub(channel)
        messages = [
            llm_pb2.ChatMessage(role="system", content="你是一个乐于解答各种问题的助手，你的任务是为用户提供专业、准确、有见地的建议。"),
            llm_pb2.ChatMessage(role="user", content=question),
        ]
        while True:
            request = llm_pb2.ChatRequest(session_id="4", messages=messages)
            print("服务器响应：")
            responses = stub.ChatStream(request)
            for response in responses:
                print(response.delta, end="", flush=True)
            print()

            # 继续对话
            next_question = input("\n您还有其他问题吗？请输入（或输入 'exit' 退出）：")
            if next_question.lower() == 'exit':
                break
            messages.append(llm_pb2.ChatMessage(role="user", content=next_question))
            
if __name__ == '__main__':
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="向 gRPC 服务器提问")
    parser.add_argument("question", type=str, help="请输入您的问题")
    args = parser.parse_args()


    # 调用 gRPC 客户端
    run(args.question)